#include <catch2/catch_test_macros.hpp>
#include "MultiGrid.h"
#include "BoundaryConditions.h"
#include "norm.h"

Real bottom(const Point<2> &p)
{
    return p[1];
}

Real up(const Point<2> &p)
{
    return 1 - p[1];
}

template <int N>
Real left(const Point<N> &p)
{
    return p[0];
}

template <int N>
Real right(const Point<N> &p)
{
    return 1 - p[0];
}

template <int C>
Real constant(const Point<2> &p)
{
    return C;
}

SingleBoundaryCondition<2> Bottom(bottom, constant<2>, 1, 0);
SingleBoundaryCondition<2> Up(up, constant<2>, 1, 0);
SingleBoundaryCondition<2> Left(left<2>, constant<2>, 1, 0);
SingleBoundaryCondition<2> Right(right<2>, constant<2>, 1, 0);

GridOp::Linear_Interp<Real, 2> Linear_Interp;
GridOp::Quadra_Interp<Real, 2> Quadra_Interp;
GridOp::Injection<Real, 2> Injection;
GridOp::FullWeight<Real, 2> FullWeight;

TEST_CASE("2d model problem for u = 2", "[DIRICHLET]")
{
    MultiGridSolver<2> MGS{Linear_Interp, FullWeight};
    MGS.define_boundary(BoundaryConditions<2>({Bottom, Up, Left, Right}));
    MGS.assignF(constant<0>);
    std::vector<Real> solution, e, norm;
    for (int n = 16; n <= 128; n *= 2)
    {
        std::cout << "Using the grid with size : " << n << std::endl;
        MGS.discretization(n);
        solution.resize(std::pow(n + 1, 2));
        e.resize(std::pow(n + 1, 2));
        MGS.solve(solution, "FMG", 2, 2);
        for (int i = 0; i <= n; ++i)
        {
            for (int j = 0; j <= n; ++j)
            {
                e[i * (n + 1) + j] = solution[i * (n + 1) + j] - 2;
            }
        }
        norm.push_back(my_vec_norm(e.data(), pow(n + 1, 2), -1));
        printf("n : %d ; max-norm : %f\n", n, norm.back());
    }

    printf("Convergence Rate : %f\n", my_conv_rate(norm.data(), 4));

    // std::cout << "The maximum norm with different grid sizes : " << std::endl;
    // for (int n = 32, i = 0; n <= 256; n *= 2, ++i)
    // {
    //     printf("n : %d ; max-norm : %f", n, norm[i]);
    // }
}

Real f1(const Point<2> &p)
{
    return 2 * ((1 - 6 * p[0] * p[0]) * (std::pow(p[1], 2) - std::pow(p[1], 4)) + (1 - 6 * p[1] * p[1]) * (std::pow(p[0], 2) - std::pow(p[0], 4)));
}

Real u1(const Point<2> &p)
{
    Real x = p[0], y = p[1];
    return (x * x - std::pow(x, 4)) * (std::pow(y, 4) - y * y);
}

TEST_CASE("2d problem u = (x^2-x^4)(y^4-y^2)", "[DIRICHLET]")
{
    Bottom.g = constant<0>;
    Up.g = constant<0>;
    Left.g = constant<0>;
    Right.g = constant<0>;
    MultiGridSolver<2> MGS{Linear_Interp, FullWeight};
    MGS.define_boundary(BoundaryConditions<2>({Bottom, Up, Left, Right}));
    MGS.assignF(f1);
    MGS.isDirectBottomSolver = true;
    std::vector<Real> solution, e, norm;

    for (int n = 16; n <= 128; n *= 2)
    {
        Real h = 1.0 / n;
        std::cout << "Using the grid with size : " << n << std::endl;
        MGS.discretization(n);
        solution.resize(std::pow(n + 1, 2));
        e.resize(std::pow(n + 1, 2));
        MGS.solve(solution, "FMG", 2, 2);
        for (int i = 0; i <= n; ++i)
        {
            for (int j = 0; j <= n; ++j)
            {
                e[i * (n + 1) + j] = solution[i * (n + 1) + j] - u1({j * h, i * h});
            }
        }
        norm.push_back(my_vec_norm(e.data(), pow(n + 1, 2), -1));
        printf("n : %d ; max-norm : %f\n", n, norm.back());
    }

    printf("Convergence Rate : %f\n", my_conv_rate(norm.data(), 4));
}

// 2-dimensional Neumann Condition is unsolvable!!!!

SingleBoundaryCondition<2> b_neum(bottom, constant<0>, 0, 1);
SingleBoundaryCondition<2> u_neum(up, constant<2>, 0, 1);
SingleBoundaryCondition<2> l_neum(left<2>, constant<0>, 0, 1);
SingleBoundaryCondition<2> r_neum(right<2>, constant<2>, 0, 1);

TEST_CASE("2d problem u = x^2+y^2", "[Neumann]")
{
    MultiGridSolver<2> MGS{Linear_Interp, FullWeight};
    MGS.define_boundary(BoundaryConditions<2>({b_neum, u_neum, l_neum, r_neum}));
    MGS.assignF(constant<-4>);
    MGS.isDirectBottomSolver = true;
    std::vector<Real> solution, e, norm;
    for (int n = 16; n <= 128; n *= 2)
    {
        Real h = 1 / n;
        std::cout << "Using the grid with size : " << n << std::endl;
        MGS.discretization(n);
        solution.resize(std::pow(n + 1, 2));
        e.resize(std::pow(n + 1, 2));
        MGS.solve(solution, "FMG", 2, 2);
        for (int i = 0; i <= n; ++i)
        {
            for (int j = 0; j <= n; ++j)
            {
                e[i * (n + 1) + j] = solution[i * (n + 1) + j] - (i * i + j * j) * h * h - solution[0];
            }
        }
        norm.push_back(my_vec_norm(e.data(), pow(n + 1, 2), -1));
        printf("n : %d ; max-norm : %f\n", n, norm.back());
        std::cout << norm.back() << std::endl;
    }
    printf("Convergence Rate : %f\n", my_conv_rate(norm.data(), 4));
}
